from ApiBase import apiBase
import re,json,os
from string import Template
import xml.etree.ElementTree as ET
from ApiModels import VectInfo,apiModels

# 基于向量库调用工具
class ApiTol():
    def __init__(self,llm,vector):
        apiBase.getLLMConnect()
        self.llm=llm
        self.vector=vector
        self.dict={}
    
    # 读取hpl工具到数据库，描述是q，函数名是a，保存到向量库里面fun_name是tool；
    def load_file(self,prefix,name:str,fname:str):
        param=name.split(".")
        if len(param) !=2:
            return
        tool_name=prefix+param[0]
        ext = param[1]
        desc=None
        # 如果固定run.hpl文件
        if ext == "hpl":
            tree = ET.parse(fname)
            root = tree.getroot()            
            element = root.find('info/description')
            desc = element.text
            if desc is None:
                desc = tool_name
            code = f'''
from ApiBase import apiBase
print('{tool_name} == >'+param)
ret=apiBase.call_pip('{fname}',param)
print(ret)
'''
        if ext == "py":
            code = f'''
def {tool_name}(param):
    from ApiTools import apiTools,apiBase
    print('{tool_name} == >'+param)
    ret=apiBase.run_python_file('{fname}',param)
    print(ret)
'''
        if ext == "sql":
            code = f'''
def {tool_name}(param):
    from ApiTools import apiTools,apiBase
    print('{tool_name} == >'+param)
    ret=apiBase.run_sql_file('{fname}',param)
    print(ret)
'''
        self.load_code(tool_name,code,desc,"py")
    
    # 读取python工具到数据库
    def load_code(self, name,code:str,remark='',type="py"):
        sql="insert into ent_tool(name,type,code,remark)values(%s,%s,%s,%s)"
        return apiBase.execute([sql],[(name,type,code,remark)],True)
       
    
    # 把目录下面的工具全部读取到表
    def load_dir(self,prefix:str,dir:str):        
        template = Template(dir)
        toolspath = template.substitute(PROJECT_HOME=os.environ["PROJECT_HOME"])
        if os.path.isfile(dir):
            name = os.path.basename(dir)
            self.load_file(prefix,name,dir)
            return 
        # 判断目录是否存在
        if not os.path.exists(toolspath) and os.path.isdir(toolspath):
            return None
        self.tools_path = toolspath
        if toolspath.endswith('/'):
            toolspath = toolspath[:-1]
        folder_names = os.listdir(toolspath)
        for name in folder_names:
            file=toolspath + "/" + name
            if os.path.isfile(file):
                self.load_file(prefix,name,file)
            else:
                self.load_dir(prefix,file) 
    
    # 根据tool的名称关键字或prompt获取工具实际的名称？？？？
    def get_tool(self,vcts:list[VectInfo],prompt,fun_name='tool-name'):        
        # 如果有工具名称，直接调用工具
        tools=[]        
        # 没有工具名称，可以通过prompt选择工具柜
        tool_names=self.vector.clts_query(vcts, prompt,fun_name=fun_name,type=0)
        if len(tool_names) == 0:
            return None
        # 只是取第1个
        name=tool_names[0]
        func_name=name['answer']
        if func_name in self.dict:
            return self.dict[func_name]
        if 'database-sql-tools' in func_name:
            uri=func_name.split("|")[1]
            # 顺序执行
            from langchain_community.tools.sql_database.tool import ListSQLDatabaseTool, InfoSQLDatabaseTool            
            from sqlalchemy import create_engine
            from langchain_community.utilities.sql_database import SQLDatabase
            engine = create_engine(uri)
            targetdb=SQLDatabase(engine)
            list_tool = ListSQLDatabaseTool(db=targetdb)
            tools.append(list_tool)
            info_tool = InfoSQLDatabaseTool(db=targetdb)
            tools.append(info_tool)                            
        else:
            tool=apiModels.get_tool(func_name)
            tools.append(tool)
        self.dict[func_name]=tools
        return tools
            
    # 大模型根据向量表找出工具的名称，并且执行，返回结果???原来的ddl接口要重新测试
    def tol_chat(self,vcts:list[VectInfo],prompt,param,fun_name='tool-name'):
        try:
            # 选择工具
            tools=self.get_tool(vcts,prompt,fun_name)
            if tools is None:
                return None
            if not param:
                param=self.vector.askQA(vcts,'You are an expert in extracting function parameters',prompt,'tool-param')
            #return tol["fun"]("postgresql://apjaimvp:mvppassword@10.220.130.88:55432/pubs3postgres,authors")
            
            # 只是运行 最相似的工具
            #ret=tools[0]['fun'].invoke("postgresql://apjaimvp:mvppassword@10.220.130.88:55432/pubs3postgres,authors")
            ret=tools[0]['fun'].invoke(param)
            return ret
        except Exception as e:
           return f'error :{str(e)}'
    
    